#include "ka.h"
#include <stdbool.h>
#include <string.h>

#include <tee_client_api.h>

#include <api_test_ta.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/bn.h>
#include <openssl/evp.h>
#include <openssl/engine.h>
#include <assert.h>

#include <sys/time.h>
static uint64_t getTimeMS()
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return (tv.tv_sec * 1000 + tv.tv_usec / 1000);
}

#define CHECK_RC(expr)                                            \
    do                                                            \
    {                                                             \
        uint32_t __rc = expr;                                     \
        if (__rc)                                                 \
        {                                                         \
            printf(__FILE__ " %d "                                \
                            " " #expr " failed with rc 0x%08x\n", \
                   __LINE__, __rc);                               \
            return __rc;                                          \
        }                                                         \
    } while (0)

void deinit_rsa();

bool inited = false;
TEEC_Context ctx;
TEEC_Session sess;
int ka_init()
{
    if (inited)
        return 0;

    uint32_t err_origin;
    TEEC_UUID uuid = TA_API_TEST_UUID;
    CHECK_RC(TEEC_InitializeContext(NULL, &ctx));
    CHECK_RC(TEEC_OpenSession(&ctx, &sess, &uuid, 0xC0000000, NULL, NULL, &err_origin));
    inited = true;
    return 0;
}

void ka_deinit()
{
    if (!inited)
        return;
    TEEC_CloseSession(&sess);

    TEEC_FinalizeContext(&ctx);
    deinit_rsa();
}

// void print_hex(uint8_t *data, size_t len)
// {
//     for (size_t i = 0; i < len; i++)
//     {
//         printf("%02x", data[i]);
//     }
//     printf("\n");
// }
#define BUFF_SIZE (RSA_SIZE / 8)
uint8_t rsa_key[8][BUFF_SIZE] = {0};
size_t rsa_key_size[8] = {0};
enum
{
    RSA_KEY_PART_N = 0,
    RSA_KEY_PART_E,
    RSA_KEY_PART_D,
    RSA_KEY_PART_P,
    RSA_KEY_PART_Q,
    RSA_KEY_PART_DMP1,
    RSA_KEY_PART_DMQ1,
    RSA_KEY_PART_IQMP,
    RSA_KEY_PART_NUM,
};
BIGNUM *rsa_key_parts[RSA_KEY_PART_NUM];

void conver_rsa_key(bool is_public)
{
    // printf("Conver RSA key to BIGNUM\n");
    int conver_num = is_public ? 2 : RSA_KEY_PART_NUM;
    for (int i = 0; i < conver_num; i++)
    {
        rsa_key_parts[i] = BN_bin2bn(rsa_key[i], rsa_key_size[i], NULL);
        // BN_print_fp(stdout, rsa_key_parts[i]);
        // printf("\n");
    }
}

RSA *rsa = NULL;
char *host_path = NULL;
bool is_self_rsa = false;
int init_rsa(bool is_public)
{
    RSA_free(rsa);
    rsa = RSA_new();
    is_self_rsa = false;
    int rc;
    if (is_public)
    {
        rc = RSA_set0_key(rsa, rsa_key_parts[RSA_KEY_PART_N], rsa_key_parts[RSA_KEY_PART_E], NULL);
        CHECK_RC(rc == 0);
        return 0;
    }
    rc = RSA_set0_key(rsa, rsa_key_parts[RSA_KEY_PART_N], rsa_key_parts[RSA_KEY_PART_E], rsa_key_parts[RSA_KEY_PART_D]);
    CHECK_RC(rc == 0);

    rc = RSA_set0_factors(rsa, rsa_key_parts[RSA_KEY_PART_P], rsa_key_parts[RSA_KEY_PART_Q]);
    CHECK_RC(rc == 0);
    rc = RSA_set0_crt_params(rsa, rsa_key_parts[RSA_KEY_PART_DMP1], rsa_key_parts[RSA_KEY_PART_DMQ1], rsa_key_parts[RSA_KEY_PART_IQMP]);
    CHECK_RC(rc == 0);
    is_self_rsa = true;
    return 0;
}

void deinit_rsa()
{
    RSA_free(rsa);
    rsa = NULL;
    host_path = NULL;
    is_self_rsa = false;
}

int host_get_rsa()
{
    host_path = NULL;
    // uint64_t start_time = getTimeMS();
    TEEC_Operation op[2] = {0};
    uint32_t err_origin;

    op[0].paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_OUTPUT, TEEC_MEMREF_TEMP_OUTPUT,
                                        TEEC_MEMREF_TEMP_OUTPUT, TEEC_MEMREF_TEMP_OUTPUT);

    op[0].params[0].tmpref.buffer = rsa_key[0];
    op[0].params[0].tmpref.size = BUFF_SIZE;
    op[0].params[1].tmpref.buffer = rsa_key[1];
    op[0].params[1].tmpref.size = BUFF_SIZE;
    op[0].params[2].tmpref.buffer = rsa_key[2];
    op[0].params[2].tmpref.size = BUFF_SIZE;
    op[0].params[3].tmpref.buffer = rsa_key[3];
    op[0].params[3].tmpref.size = BUFF_SIZE;

    CHECK_RC(TEEC_InvokeCommand(&sess, TA_API_TEST_CMD_RSA_PRI_0, &op[0], &err_origin));

    // for (int i = 0; i < 4; i++)
    // {
    //     printf("RSA[%d] size : %ld data : %d\n", i, op[0].params[i].tmpref.size, rsa_key[i][0]);
    // }

    op[1].paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_OUTPUT, TEEC_MEMREF_TEMP_OUTPUT,
                                        TEEC_MEMREF_TEMP_OUTPUT, TEEC_MEMREF_TEMP_OUTPUT);

    op[1].params[0].tmpref.buffer = rsa_key[4];
    op[1].params[0].tmpref.size = BUFF_SIZE;
    op[1].params[1].tmpref.buffer = rsa_key[5];
    op[1].params[1].tmpref.size = BUFF_SIZE;
    op[1].params[2].tmpref.buffer = rsa_key[6];
    op[1].params[2].tmpref.size = BUFF_SIZE;
    op[1].params[3].tmpref.buffer = rsa_key[7];
    op[1].params[3].tmpref.size = BUFF_SIZE / 2;
    CHECK_RC(TEEC_InvokeCommand(&sess, TA_API_TEST_CMD_RSA_PRI_1, &op[1], &err_origin));

    // for (int i = 0; i < 4; i++)
    // {
    //     printf("RSA[%d] size : %ld data : %d\n", i + 4, op[1].params[i].tmpref.size, rsa_key[i + 4][0]);
    // }

    for (size_t i = 0; i < 4; i++)
    {
        rsa_key_size[i] = op[0].params[i].tmpref.size;
        rsa_key_size[i + 4] = op[1].params[i].tmpref.size;
    }
    // for (size_t i = 0; i < 8; i++)
    // {
    //     print_hex(rsa_key[i], rsa_key_size[i]);
    // }
    conver_rsa_key(false);
    // uint64_t total_time = getTimeMS() - start_time;
    // printf("Hot Get RSA Time : %ld ms\n",total_time);
    return 0;
}

int guest_get_rsa(const char *host)
{
    char *resolved = realpath(host, NULL);
    if (resolved == NULL)
    {
        return -1;
    }
    host_path = resolved;
    // printf("Absolute host path: %s\n", resolved);

    TEEC_Operation op = {0};
    uint32_t err_origin;

    op.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT, TEEC_MEMREF_TEMP_OUTPUT,
                                     TEEC_MEMREF_TEMP_OUTPUT, TEEC_NONE);

    op.params[0].tmpref.buffer = resolved;
    op.params[0].tmpref.size = strlen(resolved);
    op.params[1].tmpref.buffer = rsa_key[0];
    op.params[1].tmpref.size = BUFF_SIZE;
    op.params[2].tmpref.buffer = rsa_key[1];
    op.params[2].tmpref.size = BUFF_SIZE;

    TEEC_Result res = TEEC_InvokeCommand(&sess, TA_API_TEST_CMD_RSA_PUB, &op, &err_origin);
    free(resolved);
    CHECK_RC(res);

    for (int i = 1; i < 3; i++)
    {
        rsa_key_size[i - 1] = op.params[i].tmpref.size;
        // printf("RSA[%d] size : %ld data : %d\n", i, op.params[i].tmpref.size, ((uint8_t *)op.params[i].tmpref.buffer)[0]);
    }
    conver_rsa_key(true);
    return 0;
}

// operate up to half RSA_SIZE at a time
// const static int operate_size = RSA_SIZE / 8 - 11; //(RSA_SIZE / 2 / 8);
const static int operate_size = (RSA_SIZE / 2 / 8);

int ka_guest_encrypt(const char *host, uint8_t *idata, int ilen, uint8_t *odata, int *olen)
{
    if (!inited)
        return -1;
    if (!host_path || strcmp(host, host_path))
    {
        CHECK_RC(guest_get_rsa(host));
        CHECK_RC(init_rsa(true));
    }
    int loop_count = ilen / operate_size + 1;
    if (*olen < loop_count * RSA_SIZE / 8)
        return -1;
    int len = 0;
    while (ilen)
    {
        assert(loop_count);
        int operate_len = ilen > operate_size ? operate_size : ilen;
        int ooperate_num = RSA_public_encrypt(operate_len, idata, odata, rsa, RSA_PKCS1_PADDING);
        ilen -= operate_len;
        idata += operate_len;
        len += ooperate_num;
        odata += ooperate_num;
        loop_count--;
    }
    *olen = len;
    // deinit_rsa();
    return 0;
}
int ka_guest_decrypt(const char *host, uint8_t *idata, int ilen, uint8_t *odata, int *olen)
{
    if (!inited)
        return -1;
    if (!host_path || strcmp(host, host_path))
    {
        CHECK_RC(guest_get_rsa(host));
        CHECK_RC(init_rsa(true));
    }
    int loop_count = ilen / (RSA_SIZE / 8);
    if (*olen < loop_count * operate_size)
        return -1;
    int len = 0;
    while (ilen)
    {
        assert(loop_count);
        int operate_len = (RSA_SIZE / 8);
        int ooperate_num = RSA_public_decrypt(operate_len, idata, odata, rsa, RSA_PKCS1_PADDING);
        ilen -= operate_len;
        idata += operate_len;
        len += ooperate_num;
        odata += ooperate_num;
        loop_count--;
    }
    *olen = len;
    // deinit_rsa();
    return 0;
}
int ka_host_encrypt(uint8_t *idata, int ilen, uint8_t *odata, int *olen)
{
    if (!inited)
        return -1;
    if (!is_self_rsa)
    {
        CHECK_RC(host_get_rsa());
        CHECK_RC(init_rsa(false));
    }
    int loop_count = ilen / operate_size + 1;
    if (*olen < loop_count * RSA_SIZE / 8)
        return -1;
    int len = 0;
    while (ilen)
    {
        assert(loop_count);
        int operate_len = ilen > operate_size ? operate_size : ilen;
        int ooperate_num = RSA_private_encrypt(operate_len, idata, odata, rsa, RSA_PKCS1_PADDING);
        ilen -= operate_len;
        idata += operate_len;
        len += ooperate_num;
        odata += ooperate_num;
        loop_count--;
    }
    *olen = len;
    // deinit_rsa();
    return 0;
}
int ka_host_decrypt(uint8_t *idata, int ilen, uint8_t *odata, int *olen)
{
    if (!inited)
        return -1;
    if (!is_self_rsa)
    {
        CHECK_RC(host_get_rsa());
        CHECK_RC(init_rsa(false));
    }
    int loop_count = ilen / (RSA_SIZE / 8);
    if (*olen < loop_count * operate_size)
        return -1;
    int len = 0;
    while (ilen)
    {
        assert(loop_count);
        int operate_len = RSA_SIZE / 8;
        int ooperate_num = RSA_private_decrypt(operate_len, idata, odata, rsa, RSA_PKCS1_PADDING);
        ilen -= operate_len;
        idata += operate_len;
        len += ooperate_num;
        odata += ooperate_num;
        loop_count--;
    }
    *olen = len;
    // deinit_rsa();
    return 0;
}

int ka_add_guest(const char *guest)
{
    if (!inited)
        return -1;
    TEEC_Result res;
    TEEC_Operation op;
    uint32_t err_origin;
    char *resolved = realpath(guest, NULL);
    if (resolved == NULL)
    {
        return -1;
    }

    memset(&op, 0, sizeof(op));

    op.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT, TEEC_NONE, TEEC_NONE, TEEC_NONE);
    op.params[0].tmpref.buffer = resolved;
    op.params[0].tmpref.size = strlen(resolved);

    res = TEEC_InvokeCommand(&sess, TA_API_TEST_CMD_ADD, &op, &err_origin);
    free(resolved);
    CHECK_RC(res);
    return 0;
}
int ka_delete_guest(const char *guest)
{
    if (!inited)
        return -1;
    TEEC_Result res;
    TEEC_Operation op;
    uint32_t err_origin;

    char *resolved = realpath(guest, NULL);
    if (resolved == NULL)
    {
        return -1;
    }
    memset(&op, 0, sizeof(op));

    op.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT, TEEC_NONE, TEEC_NONE, TEEC_NONE);
    op.params[0].tmpref.buffer = resolved;
    op.params[0].tmpref.size = strlen(resolved);

    res = TEEC_InvokeCommand(&sess, TA_API_TEST_CMD_DELETE, &op, &err_origin);
    free(resolved);
    CHECK_RC(res);
    return 0;
}